# _*_ coding: utf-8 _*_
import os
import numpy as np
import matplotlib.pyplot as plt
from copy import deepcopy
from random import sample, shuffle, uniform
from time import time, gmtime, strftime, localtime
import argparse
from collections import Counter

from VCMTSP import Instance

np.set_printoptions(suppress=True, linewidth=999, threshold=np.inf)
plt.rc('font', family='Times New Roman', weight='normal', size=12)

parser = argparse.ArgumentParser()
parser.add_argument("--case", required=True, type=str, help="The case name!")
parser.add_argument("-m", required=True, type=int, help="The number of salesmen!")
args = parser.parse_args()

m = args.m
case = args.case

instance = Instance(file=f'../../data/TSPLIB_VC/{case}.tsp/{case}.tsp.txt', m=m)
Distance = instance.distance_matrix  # distance matrix
Cost = instance.cost  # traveling cost of a single path
A = instance.A  # accessibility matrix
D = instance.depot  # depot
V = deepcopy(instance.cities[1:]).tolist()  # city set, exclude depot


def make_dir(path):
    path = path.strip()
    path = path.rstrip("\\")
    isExists = os.path.exists(path)
    if not isExists:
        os.makedirs(path)
        return True
    else:
        print(path + ' already exists.')
        return False


def normalization(x):
    if np.min(x) != np.max(x):
        return (x - np.min(x)) / (np.max(x) - np.min(x))
    else:
        return (x - np.min(x)) + 1


def find_nearest(city, roll):
    tar, min_dis = -1, 1e99
    for c in roll:
        if min_dis > Distance[city][c]:
            min_dis = Distance[city][c]
            tar = c
    assert tar != -1
    return tar, min_dis


def salesmen_of(city):
    return np.where(A[city, :] == 1)[0].tolist()


def cities_of(salesman):
    return np.where(A[:, salesman] == 1)[0].tolist()


class Solution:
    def __init__(self, init='OPT'):
        self.routes = None
        self.length = None
        self.lengths = None
        # initialize by optimal insertion
        if init == 'OPT':
            self.routes = [[D, D] for _ in range(instance.m)]  # empty routes
            self.lengths = [0 for _ in range(len(self.routes))]
            unvisited = deepcopy(V)
            while len(unvisited) > 0:
                city = sample(unvisited, 1)[0]
                unvisited.remove(city)
                self.optimal_insert(city)
        # initialize by greedy build
        elif init == 'GRD':
            self.routes = [[D] for _ in range(instance.m)]
            unvisited = deepcopy(V)
            while unvisited:
                for r in range(instance.m):
                    access_cities = cities_of(r)  # cities access to route r
                    candidates = set(unvisited) & set(access_cities)  # unvisited
                    if candidates:
                        nearest_city = find_nearest(self.routes[r][-1], roll=candidates)
                        self.routes[r].append(nearest_city)
                        unvisited.remove(nearest_city)
            for r in range(instance.m):
                self.routes[r].append(D)
        # initialize by random
        elif init == 'RND':
            self.routes = [[D] for _ in range(instance.m)]
            unvisited = deepcopy(V)
            while unvisited:
                city = sample(unvisited, 1)[0]
                unvisited.remove(city)
                salesmen = salesmen_of(city)
                r = sample(salesmen, 1)[0]
                self.routes[r].append(city)
            for r in range(instance.m):
                self.routes[r].append(D)
        else:
            print(f'Error. init={init}')
            quit(-1)
        self.calc_length()

    def calc_length(self):
        self.lengths = []
        for route in self.routes:
            self.lengths.append(Cost(route))
        self.length = sum(self.lengths)

    def optimal_insert(self, city, salesmen=None, perform=True):
        if salesmen is None:  # if salesmen not designated, search for every possible salesmen
            salesmen = salesmen_of(city)
        elif type(salesmen) is int:
            salesmen = [salesmen]
        opt_r, opt_i, opt_cost = -1, -1, 1e99
        for r in salesmen:
            for i in range(1, len(self.routes[r])):
                p, q = self.routes[r][i - 1], self.routes[r][i]
                cost = Distance[p][city] + Distance[city][q] - Distance[p][q]
                if opt_cost > cost:
                    opt_r, opt_i, opt_cost = r, i, cost
        if perform:
            self.routes[opt_r].insert(opt_i, city)
            self.lengths[opt_r] += opt_cost
            self.length = sum(self.lengths)
        return opt_cost

    def two_opt(self):
        total_improve = 0
        for r, route in enumerate(self.routes):
            improve = 0
            for i in range(1, len(route) - 1):  # exclude depot
                _j, min_diff = -1, 1e99
                for j in range(i + 1, len(route) - 1):
                    diff = Distance[route[i]][route[j + 1]] + Distance[route[i - 1]][route[j]] \
                           - Distance[route[i - 1]][route[i]] - Distance[route[j]][route[j + 1]]
                    if min_diff > diff:
                        min_diff = diff
                        _j = j
                if min_diff < 0:
                    route = route[:i] + route[_j:i - 1:-1] + route[_j + 1:]
                    self.routes[r] = route
                    improve += min_diff
            if abs(improve) < 1e-10:
                improve = 0
            self.lengths[r] = self.lengths[r] + improve
            total_improve += improve
        self.length = sum(self.lengths)
        return total_improve

    def validate(self):
        cities_in_routes = set()
        for route in self.routes:
            counter = dict(Counter(route))
            repeat = {key: value for key, value in counter.items() if value > 1 and key != 1}
            if repeat:
                print('City repeated in routes: ')
                print(self.routes)
                print(repeat)
            cities_in_routes = cities_in_routes | set(route)
        cities_in_routes.remove(D)
        if len(cities_in_routes) != len(V):
            miss = set(V) - cities_in_routes
            print(f'Invalid routes, missing {miss}')

    def edge_set(self):
        pass


def cross(pA, pB):
    child = deepcopy(pA)  # child based on pA
    remove = []  # removed cities
    for r in range(len(child.routes)):
        # remove some cities in a_routes[m]
        diff = set(child.routes[r]) - set(pB.routes[r])  # cities in pA but not in pB
        for city in diff:
            if uniform(0, 1) < Pr:  # remove by Pr
                remove.append(city)
                child.routes[r].remove(city)
    # re-insert the removed cities
    for city in remove:
        r = sample(salesmen_of(city), 1)[0]  # choose a route by preference probability
        child.optimal_insert(city, salesmen=r)

    child.calc_length()
    child.validate()
    return child


def del_and_add(solution):
    k1, k2 = sample(range(len(solution.routes)), 2)  # random select 2 routes
    compatible = set(solution.routes[k1]) & set(cities_of(k2))  # cities in k1 which can be placed in k2
    compatible.remove(D)
    if compatible:
        city = sample(list(compatible), 1)[0]  # random select one city
        solution.routes[k1].remove(city)  # remove a city in k1
        solution.optimal_insert(city, salesmen=k2)  # optimal insert city in k2
    solution.calc_length()
    solution.validate()
    return solution


def reassign(solution):
    r = sample(range(len(solution.routes)), 1)[0]  # random select a route
    removed = []
    for city in solution.routes[r]:
        if city == D:
            continue
        salesmen = salesmen_of(city)  # query accessible salesmen
        bias = 1 - 1.0 / len(salesmen)  # drop bias
        if uniform(0, 1) < (bias * Pr):  # remove cities by Pr
            solution.routes[r].remove(city)
            removed.append(city)
    shuffle(removed)
    for city in removed:
        salesmen = salesmen_of(city)
        costs = [solution.optimal_insert(city, r, perform=False) for r in salesmen]
        fit = [1.0 / (1.0 + cost) for cost in costs]
        fit_norm = np.array([_fit / sum(fit) for _fit in fit])
        r = np.random.choice(salesmen, size=1, replace=False, p=fit_norm)[0]
        solution.optimal_insert(city, salesmen=int(r), perform=True)
    solution.calc_length()
    return solution


def tbs(P1, P2):
    P = P1 + P2
    res = []
    while len(res) < NP:
        a, b = sample(P, 2)
        P.remove(a)
        P.remove(b)
        if a.length < b.length:
            res.append(a)
        else:
            res.append(b)
    return res


def top_best(P1, P2):
    P = P1 + P2
    P.sort(key=lambda _: _.length)
    return P[:NP]


problem = 'VCMTSP'
method = 'MA'
feature = 'mutate_reassign'
plt.figure(figsize=(10, 4))
max_epoch = 1
max_gen = 500
NP = 100  # size of population
Pc = 0.5  # cross prob
Pm = 0.1  # mutation prob
Pr = 0.8  # remove prob
if __name__ == '__main__':
    date = strftime("%Y-%m-%d", localtime())
    save_path = f'./results/{date}/{case}/{m}'
    make_dir(save_path)
    save_name = f'{method}_{feature}'
    start_time = time()
    fitness = []
    optimum = Solution(init='RND')
    for epoch in range(1, max_epoch + 1):
        POP = []  # population
        # The allocation preference of each city to each salesman.
        # Initialization.
        for _ in range(NP):
            POP.append(Solution(init='RND'))
        # Progress
        pBest = deepcopy(min(POP, key=lambda _: _.length))
        gBest = deepcopy(min(POP, key=lambda _: _.length))
        gBest_list = []
        pBest_list = []
        pWorst_list = []

        oBest_list = []
        cBest_list = []
        mBest_list = []

        for gen in range(1, max_gen + 1):
            # Crossover.
            POP.sort(key=lambda _: _.length)
            oBest_list.append(min(POP, key=lambda _: _.length).length)
            offspring = []
            while len(offspring) < NP:  # produce NP offspring
                if uniform(0, 1) < Pc:
                    a, b = sample(list(range(NP)), 2)
                    Child = cross(POP[a], POP[b])
                    offspring.append(Child)
            cBest_list.append(min(POP, key=lambda _: _.length).length)
            # Mutation.
            for idx in range(NP):
                if uniform(0, 1) < Pm:
                    for _ in range(m):
                        # POP[idx] = del_and_add(POP[idx])  # del & add
                        # offspring[idx] = del_and_add(offspring[idx])  # del & add

                        POP[idx] = reassign(POP[idx])
                        # offspring[idx] = reassign(offspring[idx])

                        pass
            mBest_list.append(min(POP, key=lambda _: _.length).length)
            # 2-opt exploitation.
            for _ in range(0, NP):
                POP[_].two_opt()
                offspring[_].two_opt()
                pass
            # Select.
            # POP = tbs(POP, offspring)
            POP = tbs(POP, offspring)
            # Record gBest
            pBest = deepcopy(min(POP, key=lambda _: _.length))
            if gBest.length > pBest.length:
                gBest = deepcopy(pBest)
            pBest_list.append(pBest.length)
            gBest_list.append(gBest.length)
            pWorst_list.append(max(POP, key=lambda _: _.length).length)
            print(f'problem: {problem} {case} m={m}, method: {method}, feature: {feature},'
                  f' NP={NP}, Pc={Pc}, Pm={Pm}, Pr={Pr}. '
                  f'At epoch {epoch} gen {gen}: {round(gBest.length, 2)}', end='\r')

        #     plt.clf()
        #     plt.subplot(1, 2, 1)
        #     plt.plot(range(len(gBest_list)), gBest_list, label='gBest')
        #     plt.plot(range(len(pBest_list)), pBest_list, label='pBest')
        #     plt.plot(range(len(pWorst_list)), pWorst_list, label='pWorst')
        #
        #     plt.plot(range(len(oBest_list)), oBest_list, label='oBest')
        #     plt.plot(range(len(cBest_list)), cBest_list, label='cBest')
        #     plt.plot(range(len(mBest_list)), mBest_list, label='mBest')
        #
        #     plt.legend()
        #     plt.title(f'total length: {gBest.length:.2f}')
        #     plt.subplot(1, 2, 2)
        #     instance.draw_2(plt, routes=gBest.routes)
        #     plt.tight_layout()
        #     plt.pause(0.001)
        # plt.savefig(f'{save_path}/{save_name}_curve.png')
        # plt.show()

        fitness.append(gBest_list[-1])
        if optimum.length > gBest.length:
            optimum = deepcopy(gBest)

        mode = 'w' if epoch == 1 else 'ab'
        with open(rf'{save_path}/{save_name}_fitness.txt', mode) as f:
            np.savetxt(f, np.array(gBest_list).reshape(1, len(gBest_list)))
    # save optimal route ==============================
    route = []
    for _ in optimum.routes:
        route.extend(_)
    route = np.array(route)
    with open(rf'{save_path}/{save_name}_route.txt', 'w') as f:
        np.savetxt(f, route.reshape(1, len(route)), fmt='%d')
    # print time cost =================================
    print(f'problem: {problem} {case} m={m}, method: {method}, feature: {feature},'
          f' NP={NP}, Pc={Pc}, Pm={Pm}, Pr={Pr}. '
          f'At epoch {max_epoch} gen {max_gen}: {round(np.array(fitness).mean(), 2)}')
    print(f"Time cost {strftime('%H:%M:%S', gmtime(time() - start_time))}")
